package com.kinggm.aop;

import com.google.common.util.concurrent.RateLimiter;
import com.kinggm.ano.ApiLimit;
import com.kinggm.utils.ApiLimitUtil;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import javax.annotation.Resource;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;

@Aspect
@Component
public class ApiLimitAop {

    @Resource
    private HttpServletResponse response;

    @Before("@annotation(com.kinggm.ano.ApiLimit)")
    public void limit(JoinPoint joinPoint) {
        // 获取需要被增强的方法
        Method currentMethod = getCurrentMethod(joinPoint);
        if (currentMethod == null) {
            return;
        }
        // 获取注解参数
        String key = currentMethod.getAnnotation(ApiLimit.class).key();
        double count = currentMethod.getAnnotation(ApiLimit.class).count();

        // 获取限流器
        RateLimiter rateLimiter = ApiLimitUtil.getRateLimiter(key, count);
        boolean acquire = rateLimiter.tryAcquire();
        if (acquire) {
            System.out.println("成功获取令牌");
        } else {
            // 获取请求响应
            String result = "限流了";
            System.out.println(result);

            // 处理响应 把限流提示信息写入响应体
            dealResponse(response, result);
        }
    }


    private Method getCurrentMethod(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        return signature.getMethod();
    }


    private void dealResponse(HttpServletResponse resp, String result) {
        resp.setContentType("application/json;charset=UTF-8");
        try (ServletOutputStream outputStream = resp.getOutputStream()) {
            outputStream.write(result.getBytes(StandardCharsets.UTF_8));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
